{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Infer schema directly & self-correct via CodeChat-Bison"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Imports "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from vertexai.language_models import CodeChatSession\n",
    "from vertexai.language_models import CodeChatModel\n",
    "from google.cloud import bigquery\n",
    "import logging \n",
    "import os "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Setup logging"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "logger = logging.getLogger(__name__)\n",
    "logger.setLevel(logging.DEBUG)\n",
    "logger.addHandler(logging.StreamHandler())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Setup essentials"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "SERVICE_ACCOUNT_CREDENTIALS = './../credentials/vai-key.json'\n",
    "os.environ['GOOGLE_APPLICATION_CREDENTIALS'] = SERVICE_ACCOUNT_CREDENTIALS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "PROJECT_ID = 'arun-genai-bb'\n",
    "MODEL_NAME = 'codechat-bison@latest'\n",
    "TEMPERATURE = 0 # default value = 0\n",
    "MAX_OUTPUT_TOKENS = 2048  # length of the output response | overridding the default value which is 128\n",
    "# TOP_P = 0.95  # default value\n",
    "# TOP_K = 40  # default value\n",
    "LOCATION = 'us-central1'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "DATASET = 'flight_reservations'\n",
    "TABLES = ['customers', 'flights', 'reservations', 'transactions', 'loyality_points']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = CodeChatModel.from_pretrained(MODEL_NAME)\n",
    "bq_client = bigquery.Client()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "    SELECT *\n",
      "    FROM `arun-genai-bb.flight_reservations.INFORMATION_SCHEMA.COLUMN_FIELD_PATHS`\n",
      "    WHERE table_name in (\"customers\",\"flights\",\"reservations\",\"transactions\",\"loyality_points\")\n",
      "\n"
     ]
    }
   ],
   "source": [
    "query = f\"\"\"\n",
    "    SELECT *\n",
    "    FROM `{PROJECT_ID}.{DATASET}.INFORMATION_SCHEMA.COLUMN_FIELD_PATHS`\n",
    "    WHERE table_name in ({','.join([f'\"{table}\"' for table in TABLES])})\n",
    "\"\"\"\n",
    "logger.info(query)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>table_catalog</th>\n",
       "      <th>table_schema</th>\n",
       "      <th>table_name</th>\n",
       "      <th>column_name</th>\n",
       "      <th>field_path</th>\n",
       "      <th>data_type</th>\n",
       "      <th>description</th>\n",
       "      <th>collation_name</th>\n",
       "      <th>rounding_mode</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>arun-genai-bb</td>\n",
       "      <td>flight_reservations</td>\n",
       "      <td>transactions</td>\n",
       "      <td>transaction_id</td>\n",
       "      <td>transaction_id</td>\n",
       "      <td>INT64</td>\n",
       "      <td>None</td>\n",
       "      <td>NULL</td>\n",
       "      <td>None</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>arun-genai-bb</td>\n",
       "      <td>flight_reservations</td>\n",
       "      <td>transactions</td>\n",
       "      <td>reservation_id</td>\n",
       "      <td>reservation_id</td>\n",
       "      <td>INT64</td>\n",
       "      <td>None</td>\n",
       "      <td>NULL</td>\n",
       "      <td>None</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>arun-genai-bb</td>\n",
       "      <td>flight_reservations</td>\n",
       "      <td>transactions</td>\n",
       "      <td>amount</td>\n",
       "      <td>amount</td>\n",
       "      <td>FLOAT64</td>\n",
       "      <td>None</td>\n",
       "      <td>NULL</td>\n",
       "      <td>None</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>arun-genai-bb</td>\n",
       "      <td>flight_reservations</td>\n",
       "      <td>transactions</td>\n",
       "      <td>transaction_datetime</td>\n",
       "      <td>transaction_datetime</td>\n",
       "      <td>DATETIME</td>\n",
       "      <td>None</td>\n",
       "      <td>NULL</td>\n",
       "      <td>None</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>arun-genai-bb</td>\n",
       "      <td>flight_reservations</td>\n",
       "      <td>reservations</td>\n",
       "      <td>reservation_id</td>\n",
       "      <td>reservation_id</td>\n",
       "      <td>INT64</td>\n",
       "      <td>None</td>\n",
       "      <td>NULL</td>\n",
       "      <td>None</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>arun-genai-bb</td>\n",
       "      <td>flight_reservations</td>\n",
       "      <td>reservations</td>\n",
       "      <td>customer_id</td>\n",
       "      <td>customer_id</td>\n",
       "      <td>INT64</td>\n",
       "      <td>None</td>\n",
       "      <td>NULL</td>\n",
       "      <td>None</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>arun-genai-bb</td>\n",
       "      <td>flight_reservations</td>\n",
       "      <td>reservations</td>\n",
       "      <td>flight_id</td>\n",
       "      <td>flight_id</td>\n",
       "      <td>INT64</td>\n",
       "      <td>None</td>\n",
       "      <td>NULL</td>\n",
       "      <td>None</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>arun-genai-bb</td>\n",
       "      <td>flight_reservations</td>\n",
       "      <td>reservations</td>\n",
       "      <td>reservation_datetime</td>\n",
       "      <td>reservation_datetime</td>\n",
       "      <td>DATETIME</td>\n",
       "      <td>None</td>\n",
       "      <td>NULL</td>\n",
       "      <td>None</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>arun-genai-bb</td>\n",
       "      <td>flight_reservations</td>\n",
       "      <td>reservations</td>\n",
       "      <td>status</td>\n",
       "      <td>status</td>\n",
       "      <td>STRING</td>\n",
       "      <td>None</td>\n",
       "      <td>NULL</td>\n",
       "      <td>None</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>arun-genai-bb</td>\n",
       "      <td>flight_reservations</td>\n",
       "      <td>flights</td>\n",
       "      <td>flight_id</td>\n",
       "      <td>flight_id</td>\n",
       "      <td>INT64</td>\n",
       "      <td>None</td>\n",
       "      <td>NULL</td>\n",
       "      <td>None</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10</th>\n",
       "      <td>arun-genai-bb</td>\n",
       "      <td>flight_reservations</td>\n",
       "      <td>flights</td>\n",
       "      <td>origin</td>\n",
       "      <td>origin</td>\n",
       "      <td>STRING</td>\n",
       "      <td>None</td>\n",
       "      <td>NULL</td>\n",
       "      <td>None</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>11</th>\n",
       "      <td>arun-genai-bb</td>\n",
       "      <td>flight_reservations</td>\n",
       "      <td>flights</td>\n",
       "      <td>destination</td>\n",
       "      <td>destination</td>\n",
       "      <td>STRING</td>\n",
       "      <td>None</td>\n",
       "      <td>NULL</td>\n",
       "      <td>None</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>12</th>\n",
       "      <td>arun-genai-bb</td>\n",
       "      <td>flight_reservations</td>\n",
       "      <td>flights</td>\n",
       "      <td>departure_datetime</td>\n",
       "      <td>departure_datetime</td>\n",
       "      <td>DATETIME</td>\n",
       "      <td>None</td>\n",
       "      <td>NULL</td>\n",
       "      <td>None</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>13</th>\n",
       "      <td>arun-genai-bb</td>\n",
       "      <td>flight_reservations</td>\n",
       "      <td>flights</td>\n",
       "      <td>arrival_datetime</td>\n",
       "      <td>arrival_datetime</td>\n",
       "      <td>DATETIME</td>\n",
       "      <td>None</td>\n",
       "      <td>NULL</td>\n",
       "      <td>None</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>14</th>\n",
       "      <td>arun-genai-bb</td>\n",
       "      <td>flight_reservations</td>\n",
       "      <td>flights</td>\n",
       "      <td>carrier</td>\n",
       "      <td>carrier</td>\n",
       "      <td>STRING</td>\n",
       "      <td>None</td>\n",
       "      <td>NULL</td>\n",
       "      <td>None</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>15</th>\n",
       "      <td>arun-genai-bb</td>\n",
       "      <td>flight_reservations</td>\n",
       "      <td>flights</td>\n",
       "      <td>price</td>\n",
       "      <td>price</td>\n",
       "      <td>FLOAT64</td>\n",
       "      <td>None</td>\n",
       "      <td>NULL</td>\n",
       "      <td>None</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>16</th>\n",
       "      <td>arun-genai-bb</td>\n",
       "      <td>flight_reservations</td>\n",
       "      <td>customers</td>\n",
       "      <td>customer_id</td>\n",
       "      <td>customer_id</td>\n",
       "      <td>INT64</td>\n",
       "      <td>None</td>\n",
       "      <td>NULL</td>\n",
       "      <td>None</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>17</th>\n",
       "      <td>arun-genai-bb</td>\n",
       "      <td>flight_reservations</td>\n",
       "      <td>customers</td>\n",
       "      <td>first_name</td>\n",
       "      <td>first_name</td>\n",
       "      <td>STRING</td>\n",
       "      <td>None</td>\n",
       "      <td>NULL</td>\n",
       "      <td>None</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>18</th>\n",
       "      <td>arun-genai-bb</td>\n",
       "      <td>flight_reservations</td>\n",
       "      <td>customers</td>\n",
       "      <td>last_name</td>\n",
       "      <td>last_name</td>\n",
       "      <td>STRING</td>\n",
       "      <td>None</td>\n",
       "      <td>NULL</td>\n",
       "      <td>None</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>19</th>\n",
       "      <td>arun-genai-bb</td>\n",
       "      <td>flight_reservations</td>\n",
       "      <td>customers</td>\n",
       "      <td>email</td>\n",
       "      <td>email</td>\n",
       "      <td>STRING</td>\n",
       "      <td>None</td>\n",
       "      <td>NULL</td>\n",
       "      <td>None</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>20</th>\n",
       "      <td>arun-genai-bb</td>\n",
       "      <td>flight_reservations</td>\n",
       "      <td>customers</td>\n",
       "      <td>date_of_birth</td>\n",
       "      <td>date_of_birth</td>\n",
       "      <td>DATE</td>\n",
       "      <td>None</td>\n",
       "      <td>NULL</td>\n",
       "      <td>None</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>21</th>\n",
       "      <td>arun-genai-bb</td>\n",
       "      <td>flight_reservations</td>\n",
       "      <td>customers</td>\n",
       "      <td>created_at</td>\n",
       "      <td>created_at</td>\n",
       "      <td>DATETIME</td>\n",
       "      <td>None</td>\n",
       "      <td>NULL</td>\n",
       "      <td>None</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "    table_catalog         table_schema    table_name           column_name  \\\n",
       "0   arun-genai-bb  flight_reservations  transactions        transaction_id   \n",
       "1   arun-genai-bb  flight_reservations  transactions        reservation_id   \n",
       "2   arun-genai-bb  flight_reservations  transactions                amount   \n",
       "3   arun-genai-bb  flight_reservations  transactions  transaction_datetime   \n",
       "4   arun-genai-bb  flight_reservations  reservations        reservation_id   \n",
       "5   arun-genai-bb  flight_reservations  reservations           customer_id   \n",
       "6   arun-genai-bb  flight_reservations  reservations             flight_id   \n",
       "7   arun-genai-bb  flight_reservations  reservations  reservation_datetime   \n",
       "8   arun-genai-bb  flight_reservations  reservations                status   \n",
       "9   arun-genai-bb  flight_reservations       flights             flight_id   \n",
       "10  arun-genai-bb  flight_reservations       flights                origin   \n",
       "11  arun-genai-bb  flight_reservations       flights           destination   \n",
       "12  arun-genai-bb  flight_reservations       flights    departure_datetime   \n",
       "13  arun-genai-bb  flight_reservations       flights      arrival_datetime   \n",
       "14  arun-genai-bb  flight_reservations       flights               carrier   \n",
       "15  arun-genai-bb  flight_reservations       flights                 price   \n",
       "16  arun-genai-bb  flight_reservations     customers           customer_id   \n",
       "17  arun-genai-bb  flight_reservations     customers            first_name   \n",
       "18  arun-genai-bb  flight_reservations     customers             last_name   \n",
       "19  arun-genai-bb  flight_reservations     customers                 email   \n",
       "20  arun-genai-bb  flight_reservations     customers         date_of_birth   \n",
       "21  arun-genai-bb  flight_reservations     customers            created_at   \n",
       "\n",
       "              field_path data_type description collation_name rounding_mode  \n",
       "0         transaction_id     INT64        None           NULL          None  \n",
       "1         reservation_id     INT64        None           NULL          None  \n",
       "2                 amount   FLOAT64        None           NULL          None  \n",
       "3   transaction_datetime  DATETIME        None           NULL          None  \n",
       "4         reservation_id     INT64        None           NULL          None  \n",
       "5            customer_id     INT64        None           NULL          None  \n",
       "6              flight_id     INT64        None           NULL          None  \n",
       "7   reservation_datetime  DATETIME        None           NULL          None  \n",
       "8                 status    STRING        None           NULL          None  \n",
       "9              flight_id     INT64        None           NULL          None  \n",
       "10                origin    STRING        None           NULL          None  \n",
       "11           destination    STRING        None           NULL          None  \n",
       "12    departure_datetime  DATETIME        None           NULL          None  \n",
       "13      arrival_datetime  DATETIME        None           NULL          None  \n",
       "14               carrier    STRING        None           NULL          None  \n",
       "15                 price   FLOAT64        None           NULL          None  \n",
       "16           customer_id     INT64        None           NULL          None  \n",
       "17            first_name    STRING        None           NULL          None  \n",
       "18             last_name    STRING        None           NULL          None  \n",
       "19                 email    STRING        None           NULL          None  \n",
       "20         date_of_birth      DATE        None           NULL          None  \n",
       "21            created_at  DATETIME        None           NULL          None  "
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "schema_columns = bq_client.query(query=query).to_dataframe()\n",
    "schema_columns"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "| table_catalog   | table_schema        | table_name   | column_name          | field_path           | data_type   | description   | collation_name   | rounding_mode   |\n",
      "|:----------------|:--------------------|:-------------|:---------------------|:---------------------|:------------|:--------------|:-----------------|:----------------|\n",
      "| arun-genai-bb   | flight_reservations | transactions | transaction_id       | transaction_id       | INT64       |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | transactions | reservation_id       | reservation_id       | INT64       |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | transactions | amount               | amount               | FLOAT64     |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | transactions | transaction_datetime | transaction_datetime | DATETIME    |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | reservations | reservation_id       | reservation_id       | INT64       |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | reservations | customer_id          | customer_id          | INT64       |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | reservations | flight_id            | flight_id            | INT64       |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | reservations | reservation_datetime | reservation_datetime | DATETIME    |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | reservations | status               | status               | STRING      |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | flights      | flight_id            | flight_id            | INT64       |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | flights      | origin               | origin               | STRING      |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | flights      | destination          | destination          | STRING      |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | flights      | departure_datetime   | departure_datetime   | DATETIME    |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | flights      | arrival_datetime     | arrival_datetime     | DATETIME    |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | flights      | carrier              | carrier              | STRING      |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | flights      | price                | price                | FLOAT64     |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | customers    | customer_id          | customer_id          | INT64       |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | customers    | first_name           | first_name           | STRING      |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | customers    | last_name            | last_name            | STRING      |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | customers    | email                | email                | STRING      |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | customers    | date_of_birth        | date_of_birth        | DATE        |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | customers    | created_at           | created_at           | DATETIME    |               | NULL             |                 |\n"
     ]
    }
   ],
   "source": [
    "schema_columns = schema_columns.to_markdown(index=False)\n",
    "logger.info(schema_columns)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Utility function to convert text into SQL and automatically correct the SQL query if any execution errors occur."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_and_execute_sql(prompt, max_tries=8):\n",
    "    \"\"\"\n",
    "    Generate an SQL query using the code_gen_model and execute it using bq_client.\n",
    "    \n",
    "    Args:\n",
    "    - prompt (str): Prompt to provide to the model for generating SQL.\n",
    "    - max_tries (int): Maximum number of attempts to generate and execute SQL.\n",
    "    \n",
    "    Returns:\n",
    "    - dict: A dictionary containing the successful dataframe or error messages and prompt evolution.\n",
    "    \"\"\"\n",
    "    \n",
    "    tries = 0\n",
    "    error_messages = []\n",
    "    prompts = [prompt]\n",
    "    df = None\n",
    "\n",
    "    chat_session = CodeChatSession(model=model, \n",
    "                                   temperature=TEMPERATURE, \n",
    "                                   max_output_tokens=MAX_OUTPUT_TOKENS)\n",
    "    \n",
    "    while tries < max_tries:\n",
    "        logger.info(f'ATTEMPT: {tries+1}')\n",
    "        try:\n",
    "            # Predict SQL using the model\n",
    "            response = chat_session.send_message(prompt)\n",
    "            generated_sql_query = response.text\n",
    "            generated_sql_query = '\\n'.join(generated_sql_query.split('\\n')[1:-1])\n",
    "            logger.info('-' * 50)\n",
    "            logger.info(generated_sql_query)\n",
    "            logger.info('-' * 50)\n",
    "            # Execute SQL using BigQuery client\n",
    "            df = bq_client.query(generated_sql_query).to_dataframe()\n",
    "            logger.info('SUCCEEDED')\n",
    "            return {\"dataframe\": df, \"prompts\": prompts, \"errors\": error_messages}\n",
    "        except Exception as e:\n",
    "            logger.error('FAILED')\n",
    "            # Catch the error, store the message, and try again\n",
    "            msg = str(e)\n",
    "            error_messages.append(msg)\n",
    "            # Evolve the prompt by appending the error message and asking the model to correct it\n",
    "            prompt = f\"\"\"Encountered an error: {msg}. \n",
    "To address this, please generate an alternative SQL query response that avoids this specific error. \n",
    "Follow the instructions mentioned above to remediate the error. \n",
    "\n",
    "Modify the below SQL query to resolve the issue:\n",
    "{generated_sql_query}\n",
    "\n",
    "Ensure the revised SQL query aligns precisely with the requirements outlined in the initial question.\"\"\"\n",
    "            prompts.append(prompt)\n",
    "            tries += 1\n",
    "        logger.info('=' * 100)\n",
    "\n",
    "    return {\"dataframe\": df, \"prompts\": prompts, \"errors\": error_messages}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Test Text-to-SQL scenarios"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Construct the SEED prompt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Please craft a SQL query for BigQuery that addresses the following QUESTION provided below. \n",
      "Ensure you reference the appropriate BigQuery tables and column names provided in the SCHEMA below. \n",
      "When joining tables, employ type coercion to guarantee data type consistency for the join columns. \n",
      "Additionally, the output column names should specify units where applicable.\n",
      "\n",
      "QUESTION:\n",
      "{}\n",
      "\n",
      "SCHEMA:\n",
      "{}\n",
      "\n",
      "IMPORTANT: \n",
      "Use ONLY DATETIME and DO NOT use TIMESTAMP.\n",
      "--\n",
      "Ensure your SQL query accurately defines both the start and end of the DATETIME range.\n",
      "\n"
     ]
    }
   ],
   "source": [
    "seed_prompt = \"\"\"\n",
    "Please craft a SQL query for BigQuery that addresses the following QUESTION provided below. \n",
    "Ensure you reference the appropriate BigQuery tables and column names provided in the SCHEMA below. \n",
    "When joining tables, employ type coercion to guarantee data type consistency for the join columns. \n",
    "Additionally, the output column names should specify units where applicable.\\n\n",
    "QUESTION:\n",
    "{}\\n\n",
    "SCHEMA:\n",
    "{}\\n\n",
    "IMPORTANT: \n",
    "Use ONLY DATETIME and DO NOT use TIMESTAMP.\n",
    "--\n",
    "Ensure your SQL query accurately defines both the start and end of the DATETIME range.\n",
    "\"\"\"\n",
    "logger.info(seed_prompt)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Scenario 1: Retrieve Active Reservations for a Specific Date Range"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For this scenario, you want to find all active reservations within a specific date range."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "question = \"Provide a list of all reservations from October 10th to October 15th, 2023\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Please craft a SQL query for BigQuery that addresses the following QUESTION provided below. \n",
      "Ensure you reference the appropriate BigQuery tables and column names provided in the SCHEMA below. \n",
      "When joining tables, employ type coercion to guarantee data type consistency for the join columns. \n",
      "Additionally, the output column names should specify units where applicable.\n",
      "\n",
      "QUESTION:\n",
      "Provide a list of all reservations from October 10th to October 15th, 2023\n",
      "\n",
      "SCHEMA:\n",
      "| table_catalog   | table_schema        | table_name   | column_name          | field_path           | data_type   | description   | collation_name   | rounding_mode   |\n",
      "|:----------------|:--------------------|:-------------|:---------------------|:---------------------|:------------|:--------------|:-----------------|:----------------|\n",
      "| arun-genai-bb   | flight_reservations | transactions | transaction_id       | transaction_id       | INT64       |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | transactions | reservation_id       | reservation_id       | INT64       |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | transactions | amount               | amount               | FLOAT64     |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | transactions | transaction_datetime | transaction_datetime | DATETIME    |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | reservations | reservation_id       | reservation_id       | INT64       |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | reservations | customer_id          | customer_id          | INT64       |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | reservations | flight_id            | flight_id            | INT64       |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | reservations | reservation_datetime | reservation_datetime | DATETIME    |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | reservations | status               | status               | STRING      |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | flights      | flight_id            | flight_id            | INT64       |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | flights      | origin               | origin               | STRING      |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | flights      | destination          | destination          | STRING      |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | flights      | departure_datetime   | departure_datetime   | DATETIME    |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | flights      | arrival_datetime     | arrival_datetime     | DATETIME    |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | flights      | carrier              | carrier              | STRING      |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | flights      | price                | price                | FLOAT64     |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | customers    | customer_id          | customer_id          | INT64       |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | customers    | first_name           | first_name           | STRING      |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | customers    | last_name            | last_name            | STRING      |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | customers    | email                | email                | STRING      |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | customers    | date_of_birth        | date_of_birth        | DATE        |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | customers    | created_at           | created_at           | DATETIME    |               | NULL             |                 |\n",
      "\n",
      "IMPORTANT: \n",
      "Use ONLY DATETIME and DO NOT use TIMESTAMP.\n",
      "--\n",
      "Ensure your SQL query accurately defines both the start and end of the DATETIME range.\n",
      "\n"
     ]
    }
   ],
   "source": [
    "prompt = seed_prompt.format(question, schema_columns)\n",
    "logger.info(prompt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "ATTEMPT: 1\n",
      "--------------------------------------------------\n",
      "SELECT\n",
      "  r.reservation_id,\n",
      "  r.customer_id,\n",
      "  r.flight_id,\n",
      "  r.reservation_datetime,\n",
      "  r.status,\n",
      "  f.origin,\n",
      "  f.destination,\n",
      "  f.departure_datetime,\n",
      "  f.arrival_datetime,\n",
      "  f.carrier,\n",
      "  f.price\n",
      "FROM arun-genai-bb.flight_reservations.reservations r\n",
      "JOIN arun-genai-bb.flight_reservations.flights f\n",
      "ON r.flight_id = f.flight_id\n",
      "WHERE r.reservation_datetime BETWEEN '2023-10-10 00:00:00' AND '2023-10-15 23:59:59';\n",
      "--------------------------------------------------\n",
      "SUCCEEDED\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 102 ms, sys: 14.7 ms, total: 117 ms\n",
      "Wall time: 7.63 s\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>reservation_id</th>\n",
       "      <th>customer_id</th>\n",
       "      <th>flight_id</th>\n",
       "      <th>reservation_datetime</th>\n",
       "      <th>status</th>\n",
       "      <th>origin</th>\n",
       "      <th>destination</th>\n",
       "      <th>departure_datetime</th>\n",
       "      <th>arrival_datetime</th>\n",
       "      <th>carrier</th>\n",
       "      <th>price</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>6</td>\n",
       "      <td>6</td>\n",
       "      <td>6</td>\n",
       "      <td>2023-10-10 10:00:00</td>\n",
       "      <td>Confirmed</td>\n",
       "      <td>SEA</td>\n",
       "      <td>JFK</td>\n",
       "      <td>2023-11-25 06:00:00</td>\n",
       "      <td>2023-11-25 14:30:00</td>\n",
       "      <td>United</td>\n",
       "      <td>550.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>7</td>\n",
       "      <td>6</td>\n",
       "      <td>7</td>\n",
       "      <td>2023-10-12 11:30:00</td>\n",
       "      <td>Confirmed</td>\n",
       "      <td>JFK</td>\n",
       "      <td>MIA</td>\n",
       "      <td>2023-11-27 20:00:00</td>\n",
       "      <td>2023-11-27 23:30:00</td>\n",
       "      <td>American</td>\n",
       "      <td>380.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>8</td>\n",
       "      <td>8</td>\n",
       "      <td>8</td>\n",
       "      <td>2023-10-15 13:20:00</td>\n",
       "      <td>Confirmed</td>\n",
       "      <td>MIA</td>\n",
       "      <td>JFK</td>\n",
       "      <td>2023-11-30 10:00:00</td>\n",
       "      <td>2023-11-30 13:30:00</td>\n",
       "      <td>American</td>\n",
       "      <td>380.0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   reservation_id  customer_id  flight_id reservation_datetime     status  \\\n",
       "0               6            6          6  2023-10-10 10:00:00  Confirmed   \n",
       "1               7            6          7  2023-10-12 11:30:00  Confirmed   \n",
       "2               8            8          8  2023-10-15 13:20:00  Confirmed   \n",
       "\n",
       "  origin destination  departure_datetime    arrival_datetime   carrier  price  \n",
       "0    SEA         JFK 2023-11-25 06:00:00 2023-11-25 14:30:00    United  550.0  \n",
       "1    JFK         MIA 2023-11-27 20:00:00 2023-11-27 23:30:00  American  380.0  \n",
       "2    MIA         JFK 2023-11-30 10:00:00 2023-11-30 13:30:00  American  380.0  "
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "response = generate_and_execute_sql(prompt=prompt)\n",
    "sql_output = response['dataframe']\n",
    "sql_output"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Scenario 2: Identify customers who made reservations in the past N days."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "question = \"Identify all customers who have made flight reservations within the last 7 days.\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt = seed_prompt.format(question, schema_columns)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "ATTEMPT: 1\n",
      "--------------------------------------------------\n",
      "SELECT c.customer_id,\n",
      "       c.first_name,\n",
      "       c.last_name,\n",
      "       c.email\n",
      "FROM flight_reservations.customers c\n",
      "JOIN flight_reservations.reservations r\n",
      "ON c.customer_id = r.customer_id\n",
      "WHERE r.reservation_datetime BETWEEN DATE_SUB(CURRENT_DATETIME(), INTERVAL 7 DAY) AND CURRENT_DATETIME();\n",
      "--------------------------------------------------\n",
      "SUCCEEDED\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 65.5 ms, sys: 11.7 ms, total: 77.1 ms\n",
      "Wall time: 5.28 s\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>customer_id</th>\n",
       "      <th>first_name</th>\n",
       "      <th>last_name</th>\n",
       "      <th>email</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>11</td>\n",
       "      <td>Ian</td>\n",
       "      <td>Somerhalder</td>\n",
       "      <td>ian.s@example.com</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>13</td>\n",
       "      <td>Kate</td>\n",
       "      <td>Winslet</td>\n",
       "      <td>kate.w@example.com</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>13</td>\n",
       "      <td>Kate</td>\n",
       "      <td>Winslet</td>\n",
       "      <td>kate.w@example.com</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   customer_id first_name    last_name               email\n",
       "0           11        Ian  Somerhalder   ian.s@example.com\n",
       "1           13       Kate      Winslet  kate.w@example.com\n",
       "2           13       Kate      Winslet  kate.w@example.com"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "response = generate_and_execute_sql(prompt=prompt)\n",
    "sql_output = response['dataframe']\n",
    "sql_output"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Scenario 3: Calculate Monthly Revenue\n",
    "Calculate the total revenue generated from transactions for a given month and year."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "question = \"Calculate the total revenue generated from transactions in October 2023, specifically from all reservations with a Confirmed status.\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt = seed_prompt.format(question, schema_columns)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "ATTEMPT: 1\n",
      "--------------------------------------------------\n",
      "SELECT\n",
      "  SUM(transactions.amount) AS total_revenue_usd\n",
      "FROM\n",
      "  arun-genai-bb.flight_reservations.transactions AS transactions\n",
      "JOIN\n",
      "  arun-genai-bb.flight_reservations.reservations AS reservations\n",
      "ON\n",
      "  transactions.reservation_id = reservations.reservation_id\n",
      "WHERE\n",
      "  reservations.status = 'Confirmed'\n",
      "  AND transactions.transaction_datetime BETWEEN '2023-10-01 00:00:00' AND '2023-10-31 23:59:59';\n",
      "--------------------------------------------------\n",
      "SUCCEEDED\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 61.6 ms, sys: 11.2 ms, total: 72.8 ms\n",
      "Wall time: 8.12 s\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>total_revenue_usd</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>3860.0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   total_revenue_usd\n",
       "0             3860.0"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "response = generate_and_execute_sql(prompt=prompt)\n",
    "sql_output = response['dataframe']\n",
    "sql_output"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Scenario 4: Popular Flight Times\n",
    "Identify the most popular departure hours or days for a given day or month or year."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "question = \"Determine the departure months with the highest frequency for the year 2023.\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt = seed_prompt.format(question, schema_columns)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "ATTEMPT: 1\n",
      "--------------------------------------------------\n",
      "SELECT DATE_TRUNC(transaction_datetime, MONTH) AS departure_month,\n",
      "       COUNT(*) AS num_transactions\n",
      "FROM arun-genai-bb.flight_reservations.transactions\n",
      "WHERE transaction_datetime BETWEEN '2023-01-01' AND '2023-12-31'\n",
      "GROUP BY departure_month\n",
      "ORDER BY num_transactions DESC\n",
      "LIMIT 10;\n",
      "--------------------------------------------------\n",
      "SUCCEEDED\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 61.2 ms, sys: 11 ms, total: 72.1 ms\n",
      "Wall time: 5.45 s\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>departure_month</th>\n",
       "      <th>num_transactions</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>2023-10-01</td>\n",
       "      <td>10</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>2023-11-01</td>\n",
       "      <td>6</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "  departure_month  num_transactions\n",
       "0      2023-10-01                10\n",
       "1      2023-11-01                 6"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "response = generate_and_execute_sql(prompt=prompt)\n",
    "sql_output = response['dataframe']\n",
    "sql_output"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Scenario 5: Customer Age Group\n",
    "Group customers by age brackets and count the number in each bracket."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "question = \"Group customers into five distinct age brackets and count the number of customers in each bracket.\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt = seed_prompt.format(question, schema_columns)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "ATTEMPT: 1\n",
      "--------------------------------------------------\n",
      "SELECT\n",
      "  CASE\n",
      "    WHEN date_of_birth IS NULL THEN 'Unknown'\n",
      "    WHEN date_of_birth < '1946-01-01' THEN '0-17'\n",
      "    WHEN date_of_birth < '1965-01-01' THEN '18-34'\n",
      "    WHEN date_of_birth < '1980-01-01' THEN '35-49'\n",
      "    WHEN date_of_birth < '1995-01-01' THEN '50-64'\n",
      "    ELSE '65+'\n",
      "  END AS age_bracket,\n",
      "  COUNT(customer_id) AS num_customers\n",
      "FROM flight_reservations.customers\n",
      "GROUP BY age_bracket;\n",
      "--------------------------------------------------\n",
      "SUCCEEDED\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 64.2 ms, sys: 12.1 ms, total: 76.4 ms\n",
      "Wall time: 6.22 s\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>age_bracket</th>\n",
       "      <th>num_customers</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>50-64</td>\n",
       "      <td>7</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>65+</td>\n",
       "      <td>3</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>35-49</td>\n",
       "      <td>6</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>18-34</td>\n",
       "      <td>3</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>0-17</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "  age_bracket  num_customers\n",
       "0       50-64              7\n",
       "1         65+              3\n",
       "2       35-49              6\n",
       "3       18-34              3\n",
       "4        0-17              1"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "response = generate_and_execute_sql(prompt=prompt)\n",
    "sql_output = response['dataframe']\n",
    "sql_output"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Scenario 6: Age Calculation\n",
    "Calculate the age of customers based on their date of birth and filter those who are above X years old."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "question = \"Identify and rank all customers aged 18+ who have `Confirmed` reservations for the current month, ordered by their age. Make sure to display their ages in the result.\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt = seed_prompt.format(question, schema_columns)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "ATTEMPT: 1\n",
      "--------------------------------------------------\n",
      "WITH current_month AS (\n",
      "  SELECT DATE_TRUNC(CURRENT_DATE(), MONTH) AS start_date,\n",
      "         DATE_ADD(DATE_TRUNC(CURRENT_DATE(), MONTH), INTERVAL 1 MONTH) AS end_date\n",
      ")\n",
      "\n",
      "SELECT c.customer_id,\n",
      "       c.first_name,\n",
      "       c.last_name,\n",
      "       DATEDIFF(CURRENT_DATE(), c.date_of_birth) AS age,\n",
      "       r.reservation_id,\n",
      "       r.status\n",
      "FROM current_month AS cm\n",
      "JOIN flight_reservations.reservations AS r\n",
      "ON r.reservation_datetime BETWEEN cm.start_date AND cm.end_date\n",
      "JOIN flight_reservations.customers AS c\n",
      "ON c.customer_id = r.customer_id\n",
      "WHERE c.date_of_birth < DATE_SUB(CURRENT_DATE(), INTERVAL 18 YEAR)\n",
      "AND r.status = 'Confirmed'\n",
      "ORDER BY age DESC;\n",
      "--------------------------------------------------\n",
      "FAILED\n",
      "====================================================================================================\n",
      "ATTEMPT: 2\n",
      "--------------------------------------------------\n",
      "WITH current_month AS (\n",
      "  SELECT DATE_TRUNC(CURRENT_DATE(), MONTH) AS start_date,\n",
      "         DATE_ADD(DATE_TRUNC(CURRENT_DATE(), MONTH), INTERVAL 1 MONTH) AS end_date\n",
      ")\n",
      "\n",
      "SELECT c.customer_id,\n",
      "       c.first_name,\n",
      "       c.last_name,\n",
      "       DATE_DIFF(CURRENT_DATE(), c.date_of_birth, YEAR) AS age,\n",
      "       r.reservation_id,\n",
      "       r.status\n",
      "FROM current_month AS cm\n",
      "JOIN flight_reservations.reservations AS r\n",
      "ON r.reservation_datetime BETWEEN cm.start_date AND cm.end_date\n",
      "JOIN flight_reservations.customers AS c\n",
      "ON c.customer_id = r.customer_id\n",
      "WHERE c.date_of_birth < DATE_SUB(CURRENT_DATE(), INTERVAL 18 YEAR)\n",
      "AND r.status = 'Confirmed'\n",
      "ORDER BY age DESC;\n",
      "--------------------------------------------------\n",
      "SUCCEEDED\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 79.1 ms, sys: 18.9 ms, total: 98 ms\n",
      "Wall time: 11.1 s\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>customer_id</th>\n",
       "      <th>first_name</th>\n",
       "      <th>last_name</th>\n",
       "      <th>age</th>\n",
       "      <th>reservation_id</th>\n",
       "      <th>status</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>17</td>\n",
       "      <td>Olivia</td>\n",
       "      <td>Newton</td>\n",
       "      <td>75</td>\n",
       "      <td>17</td>\n",
       "      <td>Confirmed</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>16</td>\n",
       "      <td>Nick</td>\n",
       "      <td>Fury</td>\n",
       "      <td>58</td>\n",
       "      <td>16</td>\n",
       "      <td>Confirmed</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>13</td>\n",
       "      <td>Kate</td>\n",
       "      <td>Winslet</td>\n",
       "      <td>48</td>\n",
       "      <td>14</td>\n",
       "      <td>Confirmed</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>20</td>\n",
       "      <td>Ryan</td>\n",
       "      <td>Reynolds</td>\n",
       "      <td>47</td>\n",
       "      <td>20</td>\n",
       "      <td>Confirmed</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>15</td>\n",
       "      <td>Mary</td>\n",
       "      <td>Jane</td>\n",
       "      <td>30</td>\n",
       "      <td>15</td>\n",
       "      <td>Confirmed</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>18</td>\n",
       "      <td>Peter</td>\n",
       "      <td>Parker</td>\n",
       "      <td>22</td>\n",
       "      <td>18</td>\n",
       "      <td>Confirmed</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   customer_id first_name last_name  age  reservation_id     status\n",
       "0           17     Olivia    Newton   75              17  Confirmed\n",
       "1           16       Nick      Fury   58              16  Confirmed\n",
       "2           13       Kate   Winslet   48              14  Confirmed\n",
       "3           20       Ryan  Reynolds   47              20  Confirmed\n",
       "4           15       Mary      Jane   30              15  Confirmed\n",
       "5           18      Peter    Parker   22              18  Confirmed"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "response = generate_and_execute_sql(prompt=prompt)\n",
    "sql_output = response['dataframe']\n",
    "sql_output"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Inspect the evolution of the seed prompt and how SQL query was fixed automatically by the LLM."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "==================== ATTEMPT 1 ====================\n",
      "\n",
      "Please craft a SQL query for BigQuery that addresses the following QUESTION provided below. \n",
      "Ensure you reference the appropriate BigQuery tables and column names provided in the SCHEMA below. \n",
      "When joining tables, employ type coercion to guarantee data type consistency for the join columns. \n",
      "Additionally, the output column names should specify units where applicable.\n",
      "\n",
      "QUESTION:\n",
      "Identify and rank all customers aged 18+ who have `Confirmed` reservations for the current month, ordered by their age. Make sure to display their ages in the result.\n",
      "\n",
      "SCHEMA:\n",
      "| table_catalog   | table_schema        | table_name   | column_name          | field_path           | data_type   | description   | collation_name   | rounding_mode   |\n",
      "|:----------------|:--------------------|:-------------|:---------------------|:---------------------|:------------|:--------------|:-----------------|:----------------|\n",
      "| arun-genai-bb   | flight_reservations | transactions | transaction_id       | transaction_id       | INT64       |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | transactions | reservation_id       | reservation_id       | INT64       |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | transactions | amount               | amount               | FLOAT64     |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | transactions | transaction_datetime | transaction_datetime | DATETIME    |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | reservations | reservation_id       | reservation_id       | INT64       |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | reservations | customer_id          | customer_id          | INT64       |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | reservations | flight_id            | flight_id            | INT64       |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | reservations | reservation_datetime | reservation_datetime | DATETIME    |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | reservations | status               | status               | STRING      |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | flights      | flight_id            | flight_id            | INT64       |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | flights      | origin               | origin               | STRING      |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | flights      | destination          | destination          | STRING      |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | flights      | departure_datetime   | departure_datetime   | DATETIME    |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | flights      | arrival_datetime     | arrival_datetime     | DATETIME    |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | flights      | carrier              | carrier              | STRING      |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | flights      | price                | price                | FLOAT64     |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | customers    | customer_id          | customer_id          | INT64       |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | customers    | first_name           | first_name           | STRING      |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | customers    | last_name            | last_name            | STRING      |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | customers    | email                | email                | STRING      |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | customers    | date_of_birth        | date_of_birth        | DATE        |               | NULL             |                 |\n",
      "| arun-genai-bb   | flight_reservations | customers    | created_at           | created_at           | DATETIME    |               | NULL             |                 |\n",
      "\n",
      "IMPORTANT: \n",
      "Use ONLY DATETIME and DO NOT use TIMESTAMP.\n",
      "--\n",
      "Ensure your SQL query accurately defines both the start and end of the DATETIME range.\n",
      "\n",
      "==================== ATTEMPT 2 ====================\n",
      "Encountered an error: 400 Function not found: DATEDIFF; Did you mean date_diff? at [9:8]\n",
      "\n",
      "Location: US\n",
      "Job ID: c22b27be-d740-4a5d-b02b-02ae008b85f5\n",
      ". \n",
      "To address this, please generate an alternative SQL query response that avoids this specific error. \n",
      "Follow the instructions mentioned above to remediate the error. \n",
      "\n",
      "Modify the below SQL query to resolve the issue:\n",
      "WITH current_month AS (\n",
      "  SELECT DATE_TRUNC(CURRENT_DATE(), MONTH) AS start_date,\n",
      "         DATE_ADD(DATE_TRUNC(CURRENT_DATE(), MONTH), INTERVAL 1 MONTH) AS end_date\n",
      ")\n",
      "\n",
      "SELECT c.customer_id,\n",
      "       c.first_name,\n",
      "       c.last_name,\n",
      "       DATEDIFF(CURRENT_DATE(), c.date_of_birth) AS age,\n",
      "       r.reservation_id,\n",
      "       r.status\n",
      "FROM current_month AS cm\n",
      "JOIN flight_reservations.reservations AS r\n",
      "ON r.reservation_datetime BETWEEN cm.start_date AND cm.end_date\n",
      "JOIN flight_reservations.customers AS c\n",
      "ON c.customer_id = r.customer_id\n",
      "WHERE c.date_of_birth < DATE_SUB(CURRENT_DATE(), INTERVAL 18 YEAR)\n",
      "AND r.status = 'Confirmed'\n",
      "ORDER BY age DESC;\n",
      "\n",
      "Ensure the revised SQL query aligns precisely with the requirements outlined in the initial question.\n"
     ]
    }
   ],
   "source": [
    "for i, prompt in enumerate(response['prompts']):\n",
    "    logger.info(f'==================== ATTEMPT {i+1} ====================')\n",
    "    logger.info(prompt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".LLM-Text-to-SQL-Architectures",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
